import argparse
from collections import defaultdict
from pathlib import Path
import re

parser = argparse.ArgumentParser(description = "merge xlsx")
parser.add_argument("-i", dest = "input_list", type = str, required = True)
parser.add_argument("-db", dest = "database", type = str, nargs = '?')
parser.add_argument("-o", dest = "output_file", type = str, required = True)
args = parser.parse_args()

input_file = args.input_list
database = args.database
output_file = args.output_file

pydir = Path(__file__).parent
if database:
    database = Path(database)
else:
    database = pydir / 'aro_family.xlsx'
# 读取过滤规则的文件，对基因家族输出结果进行筛选
rule1 = defaultdict(lambda: defaultdict(str))
rule = defaultdict(str)
aro_rule = pydir / 'filter_rule.xls'
with open(aro_rule) as arule:
    for line in arule:
        line = line.strip().split('\t')
        if line[1] == 'C' or line[1] == 'D':
            rule1[line[0]] = line[2]
        rule[line[0]] = line[1]
# 记录检测数
best = defaultdict(int)
# type 记录
type_num = defaultdict(int)
# 合并的注释信息取并集
anno_merge = defaultdict(str)
# type 1 加和
type1_list = ['OXA-23', 'OXA-48', 'OXA-51', 'OXA-24']
# 加和，注释信息取最大值
max_num = defaultdict(int)
# type 2 同时检出,取最大
type2_list = ['adeABC-adeRS', 'AcrAB-TolC', 'mep', 'OqxAB', 'adeIJK', 'adeFGH', 'SmeABC', 'SmeDEF']
# type 3 同时检出,取最小
type3_list = ['MexAB-OprM', 'MexCD-OprJ', 'MexEF-OprN', 'MexXY-OprM']
# MexAB-OprM 3 MexCD-OprJ 3 MexEF-OprN 3 MexXY-OprM' 3 adeABC-adeRS >=4 AcrAB-TolC 3 OqxAB 2
# 存放同时检出的class字典 以及 抗性机制字典
drug_class = defaultdict(lambda: defaultdict(str))
rm_list = defaultdict(lambda: defaultdict(str))
tax_list = defaultdict(lambda: defaultdict(str))
# 把注释信息的各列拆分放到相应字典里，取并集
def split_anno(aro_name, anno_list):
    drug_tmp, rm_tmp, tax_tmp = anno_list.split('\t')
    drug_tmp = drug_tmp.split(';')
    rm_tmp = rm_tmp.split(';')
    tax_tmp = tax_tmp.split(';')
    for key,value in enumerate(drug_tmp):
        drug_class[aro_name][value] = 1
    for key,value in enumerate(rm_tmp):
        rm_list[aro_name][value] = 1
    for key,value in enumerate(tax_tmp):
        tax_list[aro_name][value] = 1
    drug_tmp_list = []
    for key in drug_class[aro_name].keys():
        drug_tmp_list.append(key)
    rm_tmp_list = []
    for key in rm_list[aro_name].keys():
        rm_tmp_list.append(key)
    tax_tmp_list = []
    for key in tax_list[aro_name].keys():
        tax_tmp_list.append(key)
    drug_anno = '; '.join(drug_tmp_list)
    rm_anno = '; '.join(rm_tmp_list)
    tax_anno = '; '.join(tax_tmp_list)
    result = f'{drug_anno}\t{rm_anno}\t{tax_anno}'
    anno_merge[aro_name] = f'{result}'

def type_anno(aro_name, seq_num, anno_list):       
    # 如果是加和，注释信息取并集
    split_anno(aro_name, anno_list)
    if aro_name in type1_list:
        if aro_name not in max_num:
            max_num[aro_name] = seq_num
        else:
            if seq_num > max_num[aro_name]:
                max_num[aro_name] = seq_num
        best[aro_name] += seq_num
    else:
        if aro_name in type2_list or aro_name in type3_list:
            type_num[aro_name] += 1
        if aro_name not in best:
            best[aro_name] = seq_num
        else:
            if aro_name in type3_list:
                if seq_num < best[aro_name]:
                    best[aro_name] = seq_num
            elif aro_name in type2_list:
                if seq_num > best[aro_name]:
                    best[aro_name] = seq_num
            else:
                if seq_num > best[aro_name]:
                    best[aro_name] = seq_num

# 改名字清单
import pandas as pd
db = pd.read_excel(database)
change_name = defaultdict(str)
for line in range(db.shape[0]):
    change_name[db['ARO_Name'][line]]  = db['改名字'][line]
index = defaultdict(int)
sample_id1 = ''
# 覆盖度求平均
coverage_time = defaultdict(int)
coverage_num = defaultdict(int)
with open(input_file) as ifile:
    line1 = ifile.readline().strip().split('\t')
    for key1, value1 in enumerate(line1):
        index[value1] = key1
    for line in ifile:
        line = line.strip().split('\t')
        sample_id = line[index['ID']]
        sample_id1 = sample_id
        change_id = change_name[line[index['gene']]]
        seq_num = int(line[index['count']])
        drug = line[index['drug_cn']]
        rm = line[index['med_cn']]
        tax = line[index['species_cn']]
        coverage = line[index['coverage']]
        anno_list = f'{drug}\t{rm}\t{tax}'
        if change_id == 'OprM':
            type_anno('MexAB-OprM', seq_num, anno_list)
            type_anno('MexXY-OprM', seq_num, anno_list)
            coverage_time['MexAB-OprM'] += 1
            coverage_time['MexXY-OprM'] += 1
            coverage_num['MexAB-OprM'] += float(coverage)
            coverage_num['MexXY-OprM'] += float(coverage)
        else:
            type_anno(change_id, seq_num, anno_list)   
            coverage_time[change_id] += 1
            coverage_num[change_id] += float(coverage)

# 同时检出条件
def type_out(time, need_num):
    if int(time) >= int(need_num):
        return 1
    else:
        return 0

# type 4 只报一个
type4_list = ['CTX-M', 'DHA', 'MIR', 'ACT']
def type_anno2(aro_name, db):
    db = str(db).split(',')
    for i in db:
        if i in best:
            return 0
    return 1
filter_list = ['MexAB-OprM','MexXY-OprM','MexCD-OprJ','MexEF-OprN' , 'MexJk-OprM' , 'MexGHI-OpmD','MexVW-OprM']   
#type2_list = ['adeABC-adeRS', 'AcrAB-TolC', 'mep', 'OqxAB', 'adeIJK', 'adeFGH', 'SmeABC', 'SmeDEF']
with open(output_file, 'w') as ofile:
    #head = f'ARO_Name\tbest\tDrug Class\tResistance Mechanism\tNCBI Taxonomy Name\tNCBI Taxonomy ID'
    head = f'ID\tgene\tcount\tcoverage\tdrug\tmed\tspecies'
    ofile.write(f'{head}\n')
    # 按检出数排序
    L2 = sorted(best.items(), key = lambda x:x[1], reverse = True)
    for key1, value1 in enumerate(L2):
        check = 1
        key = L2[key1][0]
        value = L2[key1][1]
        if key in type2_list or key in type3_list:
            if key == 'OqxAB' or key == 'mep':
                check = type_out(type_num[key], 2)
            elif key == 'adeABC-adeRS':
                check = type_out(type_num[key], 4)
            else:
                check = type_out(type_num[key], 3)
        elif key in rule1:
            check = type_anno2(key, rule1[key])
        if check == 1:
            ratio = int(coverage_num[key] * 100 / coverage_time[key]) / 100
            if key not in rule:
                if value < 3 or ratio <= 50:
                    continue
            else:
                if rule[key] == 'B':
                    if value <= 2 and ratio <= 50:
                        continue
                elif rule[key] == 'C':
                    if value < 3 or ratio <= 50:
                        continue
            anno = anno_merge[key].split('\t')
            drug = anno[0]
            rm = anno[1]
            tax = anno[2]
            tax_num = tax.split('; ')
            # 按空格切分，物种取前面两个单词，基因也取最后一个,beta-lactamase除外，全输出
            if key.find('beta-lactamase') == -1:
                key = key.split(' ')[-1]
            for index, value2 in enumerate(tax_num):
                value3 = value2.split(' ')
                if len(value3) > 2:
                    value2 = f'{value3[0]} {value3[1]}'
                if value2 == '铜绿假单胞菌' and rm == '抗生素外排':
                    if key in filter_list:
                        ofile.write(f'{sample_id1}\t{key}\t{value}\t{ratio}\t{drug}\t{rm}\t{value2}\n')
                else:                                  
                    ofile.write(f'{sample_id1}\t{key}\t{value}\t{ratio}\t{drug}\t{rm}\t{value2}\n')
# 检出CTX-M-15，报CTX-M-15，未检出，但检出其他CTX-M，报CTX-M
# 检出DHA-1，报DHA-1，未检出，但检出其他1，未检出，但检出其他报DHA,报DHA
# 检出MIR-1，报MIR-1，未检出，报MIR   
# 检出ACT-1，报ACT-1，未检出，报ACT


